import numpy as np
from value_cal import online_value_iteration
board_rows = 12 
board_cols = 12

traj = []

def print_board(board):
    symbols = {0: ' ', 1: 'X', 2: 'O'}
    for row in board:
        print("|" + "|".join(symbols[x] for x in row) + "|")
    print("")

def check_winner(board):
    for i in range(board_rows):
        for j in range(board_cols):
            if board[i][j] == 0:
                continue
            if i > 1 and i < board_rows-2 and j > 1 and j < board_cols-2:
                if board[i-2][j-2] == board[i-1][j-1] == board[i][j] == board[i+1][j+1] == board[i+2][j+2]:
                    return board[i][j]
                if board[i-2][j+2] == board[i-1][j+1] == board[i][j] == board[i+1][j-1] == board[i+2][j-2]:
                    return board[i][j]
            if i > 1 and i < board_rows-2:
                if board[i-2][j] == board[i-1][j] == board[i][j] == board[i+1][j] == board[i+2][j]:
                    return board[i][j]
            if j > 1 and j < board_cols-2:
                if board[i][j-2] == board[i][j-1] == board[i][j] == board[i][j+1] == board[i][j+2]:
                    return board[i][j]  
    return 0

def markov_move(board, player):
    if player == 1: 
        values = np.load('values1.npy',allow_pickle=True).item()
        visits = np.load('visits1.npy',allow_pickle=True).item()
    else:
        values = np.load('values2.npy',allow_pickle=True).item()
        visits = np.load('visits2.npy',allow_pickle=True).item()
    # state_index = np.load('state_index.npy',allow_pickle=True).item()
    possible_moves = [(i, j) for i in range(board_rows) for j in range(board_cols) if board[i][j] == 0]
    if not possible_moves:
        return None
    # move = possible_moves[np.random.randint(len(possible_moves))]
    # board[move[0]][move[1]] = player
    move_value = np.zeros(len(possible_moves))
    for move_idx in range(len(possible_moves)):
        move = possible_moves[move_idx]
        board[move[0]][move[1]] = player
        if str(board) in values.keys():
            move_value[move_idx] = values[str(board)]
            if visits[str(board)] < 10000:
                print('start online calculate!')
                traj.append(board.copy())
                values, visits = online_value_iteration(traj=traj,values=values,visits=visits,which_player=player)
                traj.pop()
                move_value[move_idx] = values[str(board)]
        else:
            traj.append(board.copy())
            values, visits = online_value_iteration(traj=traj,values=values,visits=visits,which_player=player)
            traj.pop()
            move_value[move_idx] = values[str(board)]

        board[move[0]][move[1]] = 0
    move = possible_moves[np.argmax(move_value)]
    board[move[0]][move[1]] = player
    if player == 1: 
        np.save('values1.npy', values)
        np.save('visits1.npy', visits)
    else:
        np.save('values2.npy', values)
        np.save('visits2.npy', visits)
    return board

def human_move(board, player):
    row = int(input("输入落子行数0~%d:"%(board_rows-1)))
    col = int(input("输入落子列数0~%d:"%(board_cols-1)))
    while row < 0 or row > board_rows-1 or col < 0 or col > board_cols-1:
        print("超出棋盘边界，请在棋盘内落子。")
        row = int(input("输入落子行数0~%d:"%(board_rows-1)))
        col = int(input("输入落子列数0~%d:"%(board_cols-1)))
    while board[row][col] != 0:
        print("该位置已有棋子，请更换位置落子。")
        row = int(input("输入落子行数0~%d:"%(board_rows-1)))
        col = int(input("输入落子列数0~%d:"%(board_cols-1)))
    board[row][col] = player
    return board

def play_game():
    board = np.zeros((board_rows, board_cols), dtype=int)
    player = 1
    traj.append(board.copy())
    while True:
        print_board(board)
        board = markov_move(board, player)
        traj.append(board.copy())
        winner = check_winner(board)
        if winner != 0:
            print_board(board)
            print(f"玩家 {winner} 胜利!")
            break
        if np.all(board != 0):
            print_board(board)
            print("平局!")
            break
        player = 2 if player == 1 else 1

def human_game(human_player=2):
    board = np.zeros((board_rows, board_cols), dtype=int)
    player = 1
    traj.append(board.copy())
    while True:
        print_board(board)
        if player != human_player:
            board = markov_move(board, player)
        else:
            board = human_move(board, player)
        traj.append(board.copy())
        winner = check_winner(board)
        if winner != 0:
            print_board(board)
            print(f"玩家 {winner} 胜利!")
            break
        if np.all(board != 0):
            print_board(board)
            print("平局!")
            break
        player = 2 if player == 1 else 1

def human2human():
    board = np.zeros((board_rows, board_cols), dtype=int)
    player = 1
    traj.append(board.copy())
    while True:
        print_board(board)
        board = human_move(board, player)
        traj.append(board.copy())
        winner = check_winner(board)
        if winner != 0:
            print_board(board)
            print(f"玩家 {winner} 胜利!")
            break
        if np.all(board != 0):
            print_board(board)
            print("平局!")
            break
        player = 2 if player == 1 else 1

# 自对抗
# play_game()
# 人机对抗
human_game(human_player=1)
# human2human()